#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自定义车辆检测模型训练脚本
使用YOLOv8进行自定义数据集训练
"""

import os
import yaml
from ultralytics import YOLO
from pathlib import Path
import argparse

def create_dataset_config(dataset_path, train_path, val_path, test_path=None):
    """
    创建数据集配置文件
    
    Args:
        dataset_path: 数据集根目录路径
        train_path: 训练集路径（相对于dataset_path）
        val_path: 验证集路径（相对于dataset_path）
        test_path: 测试集路径（可选）
    """
    # 车辆检测的类别配置
    nc = 5  # 类别数量
    names = {
        0: 'car',      # 轿车
        1: 'truck',    # 卡车
        2: 'bus',      # 公交车
        3: 'motorcycle', # 摩托车
        4: 'bicycle'   # 自行车
    }
    
    config = {
        'path': dataset_path,
        'train': train_path,
        'val': val_path,
        'nc': nc,
        'names': names
    }
    
    # 如果有测试集，添加到配置中
    if test_path:
        config['test'] = test_path
    
    return config

def train_model(data_config_path, model_size='n', epochs=100, imgsz=640, 
                batch_size=16, device=0, project='models/custom_train'):
    """
    训练YOLO模型
    
    Args:
        data_config_path: 数据集配置YAML文件路径
        model_size: 模型大小 ('n', 's', 'm', 'l', 'x')
        epochs: 训练轮数
        imgsz: 输入图像尺寸
        batch_size: 批次大小
        device: 设备ID ('0' for GPU 0, 'cpu' for CPU)
        project: 项目保存目录
    """
    print(f"开始训练YOLOv8{model_size}模型...")
    print(f"数据集配置: {data_config_path}")
    print(f"训练参数: epochs={epochs}, imgsz={imgsz}, batch_size={batch_size}")
    
    # 加载模型（预训练模型）
    model_name = f'yolov8{model_size}.pt'
    model = YOLO(model_name)
    
    # 训练参数
    training_args = {
        'data': data_config_path,
        'epochs': epochs,
        'imgsz': imgsz,
        'batch': batch_size,
        'device': device,
        'project': project,
        'name': 'vehicle_detection',
        'save_period': 10,  # 每10轮保存一次模型
        'patience': 20,     # 早停耐心值
        'save': True,       # 保存模型
        'cache': True,      # 缓存数据
        'workers': 4,       # 数据加载线程数
        'project': project,
        'name': 'custom_vehicle_detection',
        'exist_ok': False,
        'pretrained': True,
        'verbose': True,
        'seed': 42,
        'deterministic': True,
        'single_cls': False,
        'rect': False,
        'overlap_mask': True,
        'mask_ratio': 4,
        'dropout': 0.0,
        'val': True,
        'split': 'val',
        'save_json': True,
        'save_hub': False,
        'conf': 0.001,      # 置信度阈值
        'iou': 0.7,         # NMS IoU阈值
        'max_det': 1000,    # 最大检测数
        'half': False,
        'dnn': False,
        'plots': True,      # 生成训练图表
        'source': '',
        'show': False,
        'save_frames': False,
        'save_crop': False,
        'show_labels': True,
        'show_conf': True,
        'hide_labels': False,
        'hide_conf': False,
        'line_width': 2,
    }
    
    # 开始训练
    try:
        results = model.train(**training_args)
        print("训练完成！")
        print(f"最佳模型路径: {results.save_dir}/weights/best.pt")
        print(f"最终模型路径: {results.save_dir}/weights/last.pt")
        
        # 返回训练结果
        return results
        
    except Exception as e:
        print(f"训练过程中出现错误: {e}")
        return None

def evaluate_model(model_path, data_config_path, test_images_path=None):
    """
    评估训练好的模型
    
    Args:
        model_path: 模型权重文件路径
        data_config_path: 数据集配置路径
        test_images_path: 测试图像目录路径（可选）
    """
    print(f"开始评估模型: {model_path}")
    
    # 加载训练好的模型
    model = YOLO(model_path)
    
    # 在验证集上评估
    val_results = model.val(data=data_config_path)
    
    print("验证集评估结果:")
    print(f"mAP@0.5: {val_results.box.map}")
    print(f"mAP@0.5:0.95: {val_results.box.map50}")
    
    # 如果有测试图像，进行预测
    if test_images_path and os.path.exists(test_images_path):
        print(f"在测试图像上运行推理...")
        results = model.predict(source=test_images_path, save=True, conf=0.5)
        print("测试推理完成，结果保存在runs/predict目录")
    
    return val_results

def main():
    parser = argparse.ArgumentParser(description='自定义车辆检测模型训练')
    parser.add_argument('--data_path', type=str, default='./datasets/custom_vehicle',
                       help='数据集根目录路径')
    parser.add_argument('--train_path', type=str, default='train',
                       help='训练集相对路径')
    parser.add_argument('--val_path', type=str, default='val',
                       help='验证集相对路径')
    parser.add_argument('--model_size', type=str, default='n',
                       choices=['n', 's', 'm', 'l', 'x'],
                       help='模型大小')
    parser.add_argument('--epochs', type=int, default=100,
                       help='训练轮数')
    parser.add_argument('--batch_size', type=int, default=16,
                       help='批次大小')
    parser.add_argument('--imgsz', type=int, default=640,
                       help='输入图像尺寸')
    parser.add_argument('--device', type=str, default='0',
                       help='训练设备')
    parser.add_argument('--mode', type=str, default='train',
                       choices=['train', 'evaluate'],
                       help='运行模式：train-训练，evaluate-评估')
    parser.add_argument('--model_path', type=str, default=None,
                       help='模型路径（评估模式时使用）')
    
    args = parser.parse_args()
    
    # 创建数据集配置
    dataset_config = create_dataset_config(
        args.data_path, 
        args.train_path, 
        args.val_path
    )
    
    # 保存配置文件
    config_path = os.path.join(args.data_path, 'dataset.yaml')
    with open(config_path, 'w', encoding='utf-8') as f:
        yaml.dump(dataset_config, f, default_flow_style=False, allow_unicode=True)
    
    print(f"数据集配置已保存到: {config_path}")
    
    if args.mode == 'train':
        # 训练模式
        print("开始训练...")
        results = train_model(
            data_config_path=config_path,
            model_size=args.model_size,
            epochs=args.epochs,
            imgsz=args.imgsz,
            batch_size=args.batch_size,
            device=args.device
        )
        
        if results:
            print("训练成功完成！")
            # 评估训练好的模型
            best_model_path = os.path.join(results.save_dir, 'weights', 'best.pt')
            evaluate_model(best_model_path, config_path)
        
    elif args.mode == 'evaluate':
        # 评估模式
        if not args.model_path:
            print("评估模式需要指定模型路径 --model_path")
            return
            
        evaluate_model(args.model_path, config_path)

if __name__ == '__main__':
    main()